feat(domain): benchmark all differentiable solvers in the loop#116
Draft
andrinr wants to merge 25 commits into
Draft
feat(domain): benchmark all differentiable solvers in the loop#116andrinr wants to merge 25 commits into
andrinr wants to merge 25 commits into
Conversation
andrinr
force-pushed
the
feat/ns-grid-solver-in-loop
branch
from
July 24, 2026 15:02
2b53888 to
358e231
Compare
This was referenced Jul 27, 2026
andrinr
force-pushed
the
feat/ns-grid-solver-in-loop
branch
from
July 27, 2026 12:51
7d6a683 to
04b56d8
Compare
Contributor
📊 View the full benchmark resultsNo benchmarks ran for this PR, so there is no status report. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this PR establishes
This PR adds one common 2D solver-in-the-loop neural-correction benchmark for
all six differentiable Navier–Stokes grid solvers:
It is stacked on the opt-in canonical recurrent-state contract in #121.
JAX-CFD, INS.jl, PhiFlow, and XLB carry their native checkpoint through
training, evaluation, plotting, finite-difference validation, and
differentiation. PICT and Warp-NS are velocity-complete. Both canonical
velocity and native checkpoint state are stopped in the paired stop-gradient
control.
The benchmark makes two deliberately separate measurements:
differentiation through the autoregressive solver–corrector recurrence with
the identical stopped-state recurrence.
The second comparison is the primary differentiability result. The benchmark
is therefore best read as a recurrent credit-assignment audit: does
backpropagating through several corrected solver transitions improve a policy
over local one-step supervision? It is not a claim that a VJP alone makes the
corrected solver accurate, nor that one numerical reference is universally
neutral.
Reference design
The nonlinear task is repeated against two independently discretized targets:
dt/4dt/8dt/4dt/8Both audits start from the same Fourier-prolonged continuum initial condition
and compare seeds 0 and 100 at frames 1, 8, 24, and 36. Training is rejected if
the p95 128²-versus-256² discrepancy exceeds 0.5%.
The finite-volume method uses an FFT only to solve the periodic discrete
Poisson equation. Transport and diffusion use conservative grid stencils, so
this is not a second spectral time integrator.
Common paired protocol
Every cell uses the same 32² candidate grid, viscosity, physical time step,
16 training ICs, eight held-out ICs, three paired model seeds, 200 optimizer
updates, recurrent training to
t=1.92, and evaluation tot=2.88.The learned component is the same zero-initialized Equinox periodic residual
CNN (56,098 parameters). Only the differentiable solver transition and its VJP
change between cells.
Training uses truncated backpropagation through eight consecutive correction
intervals. Each interval advances four native solver steps (
4 × 0.02 = 0.08),applies the corrector, and feeds both the corrected velocity and native solver
checkpoint into the next interval. There is no teacher forcing inside the
sampled window. The loss averages all eight corrected states and adds a
0.1-weight terminal loss on the last provisional solver state.
Held-out evaluation is fully free-running for 36 solver–corrector intervals:
after the initial condition, no reference field is fed back into the rollout.
The paired modes are:
transitions, including solver velocity and native state;
recurrent velocity and native state are stopped at every solver output.
The solver-VJP lift is the geometric paired
error_local-control / error_full-temporal-creditratio. A correction gainabove one separately means that the full-credit corrector beats the
uncorrected solver.
Nonlinear result and reference sensitivity
Absolute error and correction gain remain reference-dependent:
The native final errors move by at most 2.7% when the reference changes, but
the trained corrected errors move by as much as 40.3% (PhiFlow). PICT is the
only corrector that beats its native solver under both references. This
separates target agreement from non-convex corrector-training sensitivity:
very similar converged targets do not guarantee the same learned endpoint.
The paired solver-VJP conclusion is more robust, but not universal:
JAX-CFD, INS.jl, PhiFlow, and XLB retain a resolved full-VJP benefit under
both references. Warp-NS remains inconclusive. PICT is the important exception:
its paired VJP effect changes from significantly harmful under the spectral
target to significantly beneficial under the finite-volume target, even though
its absolute correction gain remains above one in both tasks.
Recurrent-state admission
All six solvers pass the same nonlinear and Taylor–Green recurrence gate:
Analytic and solver-specific controls
Analytic Taylor–Green control
All six solvers pass and train against the analytic Taylor–Green solution.
JAX-CFD, INS.jl, PhiFlow, PICT, and Warp-NS improve absolute error; XLB is
close to neutral. The paired VJP effect is small near this error floor, as
expected.
Solver-specific refined-reference control
Five solvers also complete a method-local task against their own 64²,
half-step target. Those absolute errors are not ranked across solvers because
the targets differ. PICT improves absolute error by 1.119×. The paired
solver-VJP benefit is resolved for JAX-CFD (37.0%), INS.jl (29.8%), XLB
(13.3%), and PICT (4.60%); Warp-NS is inconclusive.
PhiFlow is rejected before training in this control only. Its native state
closes the common-reference task, but the chaotic refined-reference audit has
long-horizon spikes: maximum coarse/fine closure residuals are 0.0101/0.0252
and reach 22.8%/70.7% of the refinement signal.
Interpretation and limits
differentiation while holding the forward recurrence and supervision fixed.
relative to local supervision; it does not by itself say that the learned
corrector beats the uncorrected solver.
The two-reference sweep tests whether the conclusions survive one independent
discretization change; it does not establish a universally fair reference.
initial-condition distribution, viscosity, resolution, and time horizon.
fail to beat the uncorrected solver. Both quantities are reported.
Implementation state
stateis differentiable in both theTesseract input and output schemas.
finite-difference checks, long-closure checks, and self-reference generation.
prolongation for resolution audits, and a gated 128²-versus-256²
space-time-convergence check.
opt-in
optimization/solver_in_loop_reference_sensitivityexperiment.Complete merged JSON, plot inputs, full-field/physics plots, GIFs, tables, and
job provenance are in the
artifact index.
Validation and offline execution
All experiments ran offline on Kander through the shared Slurm/Pyxis solver
tooling. No hosted benchmark was triggered; this draft carries
benchmark:none.1697345;1697383;1697346,1697347,1697348,1697349,1697350,1697351;1697352,1697371,1697372,1697373,1697374,1697375;1697376;index;
21db5ff: Ruff and format passed; 522 testspassed and three were skipped (
1697344);passed; hosted Python checks are reported on this PR head.
This remains a draft PR based on #121 until the canonical recurrent-state
contract lands.